package test.ground.kafka;

import com.wang.tool.trace.TraceContext;
import com.wang.tool.trace.TraceHolder;
import org.apache.kafka.clients.consumer.ConsumerInterceptor;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.header.Header;

import java.util.Map;

public class KafkaConsumerInterceptor implements ConsumerInterceptor<String, String> {



    @Override  
    public ConsumerRecords<String, String> onConsume(ConsumerRecords<String, String> records) {  
        for (ConsumerRecord<String, String> record : records) {
            Header traceHeader = record.headers().lastHeader(TraceHolder.TRACE_ID_HEADER);
            Header rpcheader = record.headers().lastHeader(TraceHolder.RPC_ID_HEADER);
            TraceContext traceContext = null;
            if (traceHeader==null) {
                traceContext = TraceContext.newTrace();
            } else {
                traceContext = TraceContext.buildCallbackTrace(new String(traceHeader.value()), rpcheader==null?null:new String(rpcheader.value()), null);
            }
            if(traceContext!=null){
                TraceHolder.set(traceContext);
            }
        }  
        return records;  
    }

    @Override
    public void onCommit(Map<TopicPartition, OffsetAndMetadata> map) {
    }

    @Override
    public void close() {
    }

    @Override  
    public void configure(Map<String, ?> configs) { }  
}